import pathlib

import matplotlib.pyplot
import torch

import torch_plus
from simple_pendulum_dataset import SimplePendulumDataset
from simple_pendulum_problem import SimplePendulumProblem


def _save_data(directory: pathlib.Path,
               title: str,
               problem: SimplePendulumProblem,
               t_batch: torch.Tensor) -> None:
    dataset = SimplePendulumDataset(problem, t_batch)
    dataset.save(directory / f"{title}.pt")
    matplotlib.pyplot.scatter(t_batch, dataset.theta())
    matplotlib.pyplot.title(title)
    matplotlib.pyplot.xlabel(r"$t$")
    matplotlib.pyplot.ylabel(r"$\theta$")
    matplotlib.pyplot.savefig(directory / f"{title}.png")
    matplotlib.pyplot.close()


def main():
    output_directory = pathlib.Path("./data").resolve()

    random = torch.Generator()
    random.manual_seed(1234)

    _save_data(output_directory, "not_damped",
               SimplePendulumProblem(10, 0, 0.5),
               torch_plus.rand(20, 3, 10, random))

    _save_data(output_directory, "under_damping",
               SimplePendulumProblem(15, 1, 1),
               torch_plus.rand(20, 3, 10, random))

    # 在过阻尼的情况下， 3 到 10 的数据可能特征不够明显，难以得到正确结果，因此取 0 到 10
    _save_data(output_directory, "over_damping",
               SimplePendulumProblem(5, 10, 1.5),
               torch_plus.rand(20, 0, 10, random))

    print(f"The generated data have been saved to {output_directory}")


if __name__ == "__main__":
    main()
